
import numpy as np


def positionFeature(labels, Nth):
    
    tmp = np.where(labels == Nth)
    
    return  np.mean(tmp[0]),np.std(tmp[0]),np.mean(tmp[1]),np.std(tmp[1])    

def colorFeature(L,labels, Nth):
    
    tmp = np.where(labels == Nth, L)
    return np.mean(tmp),np.std(tmp)

def features( L,labels, Nth ):
    
    tmp = np.where(labels == Nth)
    
    x_pos = tmp[0]
    y_pos = tmp[1]
    col   = L[np.where(labels == Nth)] 
    
    return np.cov(np.vstack((x_pos,y_pos,col))),np.mean(tmp[0]),np.mean(tmp[1]),np.mean(col)

def featuresRGB( R, G, B, labels, Nth ):
    
    tmp = np.where(labels == Nth)
    
    x_pos = np.array(tmp[0],dtype=np.double)
    
    y_pos = np.array(tmp[1],dtype=np.double)
    
    _R = R[tmp]
    _G = G[tmp]
    _B = B[tmp]
    
    return np.cov(np.vstack((x_pos,y_pos,_R,_G,_B))),np.mean(x_pos),np.mean(y_pos),np.mean(_R),np.mean(_G),np.mean(_B)

def proba_dl_thetam(dl, miu, covarMatrix):
    
    V = np.matrix(dl - np.matrix(miu))
    exp = V* np.matrix(np.linalg.inv(covarMatrix))*V.T
    
    return np.exp(-0.5*exp)/(np.power(2*np.pi,5)*np.linalg.det(covarMatrix)) 


def prob_m_dl(m,pi_m, dl, miuArray,covarMatrixArray):
    
    tmp1 = pi_m * proba_dl_thetam(dl,miuArray[m],covarMatrixArray[m])
    
    if m == 0:
        tmp2 = pi_m * proba_dl_thetam(dl,miuArray[0],covarMatrixArray[0]) + (1-pi_m) * proba_dl_thetam(dl,miuArray[1],covarMatrixArray[1])
    elif m == 1:
        tmp2 = (1-pi_m) * proba_dl_thetam(dl,miuArray[0],covarMatrixArray[0]) + pi_m * proba_dl_thetam(dl,miuArray[1],covarMatrixArray[1])
    
    return float(tmp1/tmp2)

def pi_m__t_plus_1(R,G,B,m,pi_m,miuArray,covarMatrixArray):
    
    pixels_count = R.shape[0]*R.shape[1]
    sum = 0
    for x in xrange(R.shape[0]):
        for y in xrange(R.shape[1]):
            
            dl = np.matrix([x,y,R[x][y],G[x][y],B[x][y]])
            sum = sum + prob_m_dl(m,pi_m[m], dl, miuArray,covarMatrixArray)
    
    return sum/pixels_count

def miu_m__t_plus_1(R,G,B,m,pi_m,miuArray,covarMatrixArray):
    
    sum1 = np.matrix([0,0,0,0,0])
    sum2 =0
    for x in xrange(R.shape[0]):
        for y in xrange(R.shape[1]):
            
            dl = np.matrix([x,y,R[x][y],G[x][y],B[x][y]])
            tmp = prob_m_dl(m,pi_m, dl, miuArray,covarMatrixArray)
            
            sum1 = sum1 + dl * tmp
            sum2 = sum2 + tmp
    
    return sum1/sum2


def covarMatrix_m__t_plus_1(R,G,B,m,pi_m,miuArray,covarMatrix):
    
    sum1 = np.matrix(np.zeros((5,5)))
    sum2 =0
    for x in xrange(R.shape[0]):
        for y in xrange(R.shape[1]):
            
            dl = np.matrix([x,y,R[x][y],G[x][y],B[x][y]])
            miu_m = np.matrix(miuArray)
            tmp = prob_m_dl(m,pi_m, dl, miuArray,covarMatrix)
            sum1 = sum1 + (dl - miu_m).T * (dl - miu_m) * tmp 
            sum2 = sum2 + tmp
    
    return sum1/sum2